import os
from langchain.chat_models import init_chat_model

from app.aichat.internal.config import config
from common.rpc import stub
from common.rpc.assistant import assistant_pb2_grpc

def _set_llm_api_key(c: config.Config):
    if not os.getenv(c.llm_conf.api_key_name):
        os.environ[c.llm_conf.api_key_name] = c.llm_conf.api_key
        
def _new_llm_model(c: config.Config):
    _set_llm_api_key(c)
    return init_chat_model(c.llm_conf.model, model_provider=c.llm_conf.model_provider)

class ServiceContext:
    def __init__(self, c: config.Config):
        self.config = c

        self.llm_model = _new_llm_model(c)
        
        self.assistant_stub: assistant_pb2_grpc.AssistantStub = \
            stub.new_rpc_stub(c.assistant_rpc_conf, assistant_pb2_grpc.AssistantStub)
